function VAST_preproc_TFR(filename,STARTSTEP,loadMusc)

% Function for preprocessing VAST data for time-frequency analysis
% INPUTS ---
% filename: of the form 'p0246_a.bdf'
% STARTSTEP: where do you want to start the analysis (see below)
% loadMusc: look for manually detected artifacts from ERP pipeline?
%
% Guide for STARTSTEP: At what preprocessing stage are you starting?
% Note: If STARTSTEP > 0, necessary preceding variables will be loaded by 
% this pipeline's load_analyzed_vars function, load_analyzed_vars_VAST_TFR.
%
% LIST OF STEPS:
% Step 0: Read data channel-by-channel, immediately downsample
% Step 1: Filter between 0.5 and 50 Hz
% Step 2: Manually identify muscle artifacts (if loadMusc == 0), or
%         load artifact segments from previous ERP analysis
% Step 3: ICA to remove blinks/saccades
% Step 4: Pull relevant window triggers, save trl definition array
% Step 5: Subtract condition-and-participant-specific ERPs
% Step 6: Epoch into FULL TRIALS (with padding)
% Step 7: Peak-to-peak artifact detection in individual stim windows
% Step 8: Short-time Fourier Transform
% Step 9: Split TFR into responses in each stim window (ENC, INT, RET)


%% User defined variables

settings.do_ica = 1; % run ICA for blinks, or skip it
settings.p2pThresh_min = -100; % uV, trials that dip below this value will be rejected 
settings.p2pThresh_max = 100; % uV, trials that exceed this value will be rejected
settings.ds2 = 256; % Hz, data will be downsampled to this SR
settings.origFs = 2048;
settings.Ntrl = 240; % how many trials in the experimental session?
settings.loadMusc = loadMusc;


compname = getenv('computername');

% Check for necessary inputs
if nargin < 3
    loadMusc = 0; % define artifacts manually by default
end
if ~ismember(loadMusc,[0 1])
    error('loadMusc must be either 0 or 1')
end

if nargin < 2
    STARTSTEP = 0;
end
if STARTSTEP < 0 
    STARTSTEP = 0;
elseif STARTSTEP > 9
    error('Specified STARTSTEP exceeds the number of analysis steps')
end

if nargin < 1
    filename = 'p1556_a.bdf';
end

settings.filename = filename;
settings.filenameNoExt = settings.filename(1:end-4);

% Set Paths
% NOTE: New raw data paths imply that data are in pXXXX folders
if strcmp(compname,'DESKTOP-BHR0AU7')
    settings.filedirRaw = strcat('E:\ANL\Experiments\DATA\VAST\RAW\',filename(1:end-6),'\');
    settings.filedirAnalyzed = strcat('E:\ANL\Experiments\DATA\VAST\PREPROCESSED\TFR\',...
        settings.filenameNoExt, '\');
    erpdir = strcat('E:\ANL\Experiments\DATA\VAST\PREPROCESSED\ERP\',...
        settings.filenameNoExt, '\');
elseif strcmp(compname,'MSI')
    % Now using external SSD for laptop analysis -- should be drive E:
    settings.filedirRaw = strcat('E:\ANL\RAW_DATA\VAST\',filename(1:end-6),'\');
    settings.filedirAnalyzed = strcat('E:\ANL\PREPROCESSED_DATA\VAST\TFR\',...
        settings.filenameNoExt, '\');
    erpdir = strcat('E:\ANL\PREPROCESSED_DATA\VAST\ERP\',...
        settings.filenameNoExt, '\');
elseif strcmp(compname,'RKC-PER-WD-0002')
    settings.filedirRaw = strcat('C:\Users\mnjoroge\Documents\VAST\VAST_Raw_Data\',filename(1:end-6),'\');
    settings.filedirAnalyzed = strcat('C:\Users\mnjoroge\Documents\VAST\VAST_Preprocessed_Data\TFR\',...
        settings.filenameNoExt, '\');
    erpdir = strcat('C:\Users\mnjoroge\Documents\VAST\VAST_Preprocessed_Data\ERP\',...
        settings.filenameNoExt, '\');
else
    error('Specify paths appropriately for this machine -- then delete this error')
    settings.filedirRaw = 'U:\eng_research_anl\Justin\RESULTS\binding_popout\RAW\2stream\';
    settings.filedirAnalyzed = strcat('C:\ANL_Experiments\RESULTS\binding_popout_2stream\Preprocessed_Data\',...
        settings.filenameNoExt, '\');
end
    
if ~exist(settings.filedirAnalyzed, 'dir') % does the Analysis folder not exist?
    mkdir(settings.filedirAnalyzed) % create it if not
end

% Load any previous data and accessory variables
if STARTSTEP > 0
    load_analyzed_vars_VAST_TFR(STARTSTEP, settings.filedirAnalyzed,...
        settings.filenameNoExt, erpdir);
end

settings.wholefile = strcat(settings.filedirRaw,settings.filename);
save([settings.filedirAnalyzed settings.filenameNoExt '_settings'],...
    'settings','-v7.3');
% Add the current directory containing this function to the path
addpath(genpath(pwd));
% Add FieldTrip toolbox to the path as well
if strcmp(compname,'DESKTOP-BHR0AU7')
    addpath('C:\Users\jtjus\Documents\MATLAB\fieldtrip-20180512\fieldtrip-20180512');
    addpath(genpath('E:\MATLAB\CSDtoolbox'));
elseif strcmp(compname,'MSI')
    addpath('D:\MATLAB\fieldtrip-20180809\fieldtrip-20180809');
    addpath(genpath('D:\MATLAB\CSDtoolbox'));
elseif strcmp(compname,'RKC-PER-WD-0002')
    addpath('C:\Users\mnjoroge\Documents\MATLAB\fieldtrip-20180512\fieldtrip-20180512')
    addpath(genpath('C:\Users\mnjoroge\Documents\MATLAB\CSDtoolbox'))
else
    error('Set FieldTrip path -- then delete this error.')
end
ft_defaults % initialize minimum required FT paths

% Get Biosemi 64 channel locations
cfg = [];
cfg.layout = 'biosemi64.lay'; % layout file located in 
                              % FieldTrip folder > template > layout
q = ft_prepare_layout(cfg);


%% STEP 0: read in continuous raw data, immediately downsample
% Data will be read in one channel at a time, immediately downsampled for
% feasability

if STARTSTEP == 0
    
    % If this data has been loaded and downsampled previously for ERP
    % analysis, no need to create a new file
    ERP_step1file = strcat(erpdir,settings.filenameNoExt,'_0_ds.mat');
    if exist(ERP_step1file, 'file') == 2
        load(ERP_step1file, 'rawData');
    else
        % Settings for file read
        cfg = [];
        cfg.dataset = settings.wholefile;
        cfg.reref = 'no';
        cfg.continuous = 'yes';
        % Settings for ds
        cfg_ds = [];
        cfg_ds.resamplefs = settings.ds2;

        % Due to dataset size, read in one channel at a time...
        % Group 1) datafiles with ALL channels accidentally recorded
        if strcmp(filename, 'p0246_v.bdf')
            chans2read = [65:128, 257:264, 272];
            group = 1;
        % Group 2) datafiles with only A+B recorded
        elseif ismember(filename, {'p2317_a.bdf', 'p2317_v.bdf'})
            chans2read = 1:73;
            group = 2;
        % Group 3) datafiles with A:D recorded, due to broken A+V inputs
        else
            group = 3;
            chans2read = 65:137;
        end

        data_collected = cell(length(chans2read), 1);
        for ch = 1:length(chans2read)
            cfg.channel = chans2read(ch);
            temp = ft_preprocessing(cfg);
            % and immediately downsample
            data_collected{ch} = ft_resampledata(cfg_ds, temp);
        end

        % Properly concatenate all the individual channel data
        cfg = [];
        rawData = ft_appenddata(cfg, data_collected{:});

        % Clear unnecessary variables
        clear temp data_collected

        % Now can re-reference the data
        cfg = [];
        cfg.reref = 'yes';
        cfg.continuous = 'yes';
        cfg.refchannel = {'EXG1', 'EXG2'};
        rawData = ft_preprocessing(cfg, rawData);

        % Get the channel labels right
        if group ~= 2
            if ~(sum(contains(rawData.label, 'C32'))==1 &&...
                    sum(contains(rawData.label, 'D32'))==1)
                error("chans2read set up for broken A+B inputs, but data doesn't match this.")
            end
        end
        rawData.label(1:64) = q.label(1:64); % use actual labels instead of C+D

        % Make the eye movement data bipolar
        rawData.trial{1}(67,:) = rawData.trial{1}(67,:) - rawData.trial{1}(68,:);
        rawData.trial{1}(69,:) = rawData.trial{1}(69,:) - rawData.trial{1}(70,:);

        % Remove mastoid channels, unused EXG channels
        % Note, hard-coded for our 64 channel setup
        cfg = [];
        cfg.channel = rawData.label([1:64,67,69]);

        rawData = ft_preprocessing(cfg, rawData);

        disp(['Saving: ',settings.filedirAnalyzed, settings.filenameNoExt, '_0_raw'])
        save([settings.filedirAnalyzed settings.filenameNoExt '_0_ds'],'rawData','-v7.3');
    end % original raw data load
    
end % STEP 0


%% STEP 1: Filter data for time-frequency analysis (between 0.5 and 50 Hz)
% Will be accomplished by separate high- and low-pass filters, so
% transition widths can be set separately

if STARTSTEP <= 1

    cfg = [];
    cfg.continuous = 'yes';
    % High-pass filter
    cfg.hpfilter = 'yes';
    cfg.hpfilttype = 'firws';
    cfg.hpfreq = 0.5; % Hz
    cfg.hpfiltdf = 0.5*cfg.hpfreq;
    cfg.hpfiltwintype = 'kaiser';
    % Low-pass filter
    cfg.lpfilter = 'yes';
    cfg.lpfilttype = 'firws';
    cfg.lpfreq = 82; % Hz
    cfg.lpfiltdf = 0.15*cfg.lpfreq;
    cfg.lpfiltwintype = 'kaiser';

    rawData = ft_preprocessing(cfg,rawData);

    disp(['Saving: ',settings.filedirAnalyzed, settings.filenameNoExt, '_1_ds_filt'])
    save([settings.filedirAnalyzed settings.filenameNoExt '_1_ds_filt'], 'rawData','-v7.3');

end % Step 1


%% Step 2: Detect artifact segments for ICA

if STARTSTEP <= 2 && loadMusc == 1
    % Try loading the muscle artifacts from the ERP directory 
    ERP_step2file = strcat(erpdir,settings.filenameNoExt,'_STEP2_vars.mat');
    if exist(ERP_step2file, 'file') == 2
        load(ERP_step2file, 'badChans', 'muscArtfct');
        clear VARS
        disp('Visually-definied artifacts loaded from ERP analysis.')
        % Save here to facilitate loading later
        save(strcat(settings.filedirAnalyzed, settings.filenameNoExt,...
            '_STEP2_vars.mat'), 'badChans', 'muscArtfct');
    else % trying to load a file that does not exist
        warning('Manual artifact load from ERP data FAILED.')
        warning('Entering manual artifact marking mode...')
        loadMusc = 0;
    end
end
    
if STARTSTEP <= 2 && loadMusc == 0 % Identify the muscle artifacts by hand
    disp('Mark muscle artifacts and take note of any bad channel labels.')
    input('Press ENTER to continue')
    
    cfg = [];
    cfg.viewmode = 'vertical';
    cfg.blocksize = 30;
    cfg.ylim = [-30 30];
   
    cfg = ft_databrowser(cfg, rawData);

    % Save a note of bad channels
    temp = inputdlg('Enter SPACE-SEPARATED bad channel LABELS:',...
        'Any bad channels to interpolate?', [1 50]);
    temp = char(temp);
    badChans = strsplit(temp,' ');
    clear temp

    % SAVE 
    muscArtfct = cfg.artfctdef.visual.artifact;
    save([settings.filedirAnalyzed settings.filenameNoExt '_STEP2_vars.mat'],...
         'badChans','muscArtfct');
 
end % Step 2


%% STEP 3: ICA to remove blinks/saccades

if STARTSTEP <= 3
    
    % TEMPORARILY reject the identified muscle artifacts
    % Note: This is done for ICA cleanliness. Good ICA components will be 
    % back-projected onto the previous version of rawData.
    cfg = [];
    cfg.artfctdef.reject = 'partial';
    cfg.artfctdef.muscle.artifact = muscArtfct;
    rawData = ft_rejectartifact(cfg, rawData);
    
    % Select good channels and run the ICA
    cfg = [];
    allChans = rawData.label;
    isChanOk = cell2mat(cellfun(@(x) ~ismember(x,badChans), allChans,...
        'UniformOutput', false));
    cfg.channel = rawData.label(isChanOk);
    cfg.method = 'runica';

    components = ft_componentanalysis(cfg,rawData);
  
    % get rid of the dummy data -- will reload the unaltered data afterward
    clear rawData
    
    % Plot component topographies
    cfg = [];
    cfg.layout = 'biosemi64.lay';
    cfg.comment = 'no';
    % Fig 1: identified components 1:32
    cfg.component = 1:32;
    figure();
    ft_topoplotIC(cfg, components)
    % Fig 2: identified components 33:64
    cfg.component = 33:max(size(components.topo));
    figure();
    ft_topoplotIC(cfg, components)
    
    % Pause here to let the user check out the component topographies
    input('Press ENTER to continue to Trial View')
    
    close all
    
    % Visualize components as time series
    cfg = [];
    cfg.layout = 'biosemi64.lay';
    cfg.viewmode = 'component';

    ft_databrowser(cfg, components)
    
    % Pause here to let the user check out the component time series
    input('Press ENTER to continue, specify blink components.')
    
    % Gather the identified components
    temp = inputdlg('Enter space-separated component numbers:',...
                 'Identify Blink Components', [1 50]);
    blinkCmp = str2num(temp{:});
    clear temp

    % Remove the bad components and backproject onto the unaltered data
    load(strcat(settings.filedirAnalyzed, settings.filenameNoExt,...
        '_1_ds_filt.mat'), 'rawData');

    cfg = [];
    cfg.component = blinkCmp;
    rawData = ft_rejectcomponent(cfg, components, rawData);
    
    % ------------
    % Do some additional cleanup at this stage
    % ------------
    
    % Remove the EXG channels -- they won't be used further
    cfg = [];
    cfg.channel = rawData.label(~strcmp(rawData.label,'EXG3') & ~strcmp(rawData.label,'EXG5'));
    rawData = ft_preprocessing(cfg,rawData);
    
    % Prepare the neighborhood structure for interpolating any bad channels
    cfg = [];
    cfg.method = 'distance';
    cfg.layout = 'biosemi64.lay';   
    neighbours = ft_prepare_neighbours(cfg,rawData);
    
    % Interpolate bad channels
    cfg = [];
    cfg.method = 'spline';
    cfg.badchannel = badChans';
    cfg.neighbours = neighbours;
    cfg.layout = 'biosemi64.lay';
    % NOTE: apparently best practice is to load individualized electrode
    % positions from the polhemus data. Incorporate this down the line. See
    % documentation for ft_channelrepair and ft_read_sens
    rawData = ft_channelrepair(cfg,rawData);

    save([settings.filedirAnalyzed settings.filenameNoExt '_STEP3_vars'],...
        'components','neighbours','-v7.3');
    disp(['Saving: ',settings.filedirAnalyzed, settings.filenameNoExt, '_3_ds_filt_ica'])
    save([settings.filedirAnalyzed settings.filenameNoExt '_3_ds_filt_ica'],'rawData','-v7.3');
    
end % Step 3


%% Step 4: Pull relevant window triggers, save trl definition array

if STARTSTEP <= 4
    
    % Read events from bdf file, format them
    eventTrigs = ft_read_event(settings.wholefile);
    eVals = extractfield(eventTrigs,'value');
    eSamps = extractfield(eventTrigs,'sample');
    eTypes = extractfield(eventTrigs,'type');
    realEvents = strcmp(eTypes,'STATUS');
    if sum(realEvents) == length(eVals) % under normal conditions should always be true
        eSamps = eSamps(realEvents);
    end

    % Downsample triggers
    eSamps = round((settings.ds2 * eSamps) ./ settings.origFs);
    
    % Determine whether this is the old or new version, set triggers
    % accordingly.
    uniqTrigs = bitand(unique(eVals),255);
    if any(ismember([41 42 43],uniqTrigs) == false) % window start trigs are absent
        prompt = strcat('Window start trigs (41,42,43) not present.',...
            'Quit or reconstruct?',...
            '(Press Q to quit, anything else continues.)');
        toCont = inputdlg(prompt);
        if strcmpi(toCont,'Q')
            error('Preprocessing stopped by user.')
        else
            % Repair broken STIM1 and STIM2 triggers (10:13, 20:23)
            if ~isequal(sum(eVals==65290),sum(eVals==65291),sum(eVals==65292),...
                    sum(eVals==65293),sum(eVals==65300),sum(eVals==65301),...
                    sum(eVals==65302),sum(eVals==65303))

                warning('Data is missing some trial triggers... reconstructing them from existing triggers.')
                [eVals,eSamps] = repair_broken_trigs(eSamps,eVals,settings);
            end

            % NOTE: these windowTrigs will be adjusted in trialfun to
            % appear where the actual window start trigs occur in later
            % versions of the experiment code. See ft_trialfun_VAST_TFR for
            % more information.
            windowTrigs.stim1 = eSamps(eVals == 65290); % re. first STIM1 (10)
            windowTrigs.dist = eSamps(eVals == 65293); % re. last STIM1 (13)
            windowTrigs.stim2 = eSamps(eVals == 65300); % re. first STIM2 (20)
            windowTrigs.stim2end = eSamps(eVals == 65303); % last STIM2 (23)

            % special correction for one dataset with extra trials
            if strcmp(settings.filenameNoExt, 'p6554_a')
                windowTrigs.stim1 = windowTrigs.stim1(21:end);
                windowTrigs.dist = windowTrigs.dist(21:end);
                windowTrigs.stim2 = windowTrigs.stim2(21:end);
                windowTrigs.stim2end = windowTrigs.stim2end(21:end);
            end

            % Define the trial structure -- reconstructing windows
            [trl,windowTrigs] = ft_trialfun_VAST_TFR(windowTrigs,settings, 1); 
        end

    else % newer version: window start trigs are present
        % Repair broken STIM1 and STIM2 triggers (10:13, 20:23)
        if ~isequal(sum(eVals==65290),sum(eVals==65291),sum(eVals==65292),...
                sum(eVals==65293),sum(eVals==65300),sum(eVals==65301),...
                sum(eVals==65302),sum(eVals==65303))

            warning('Data is missing some trial triggers... reconstructing them from existing triggers.')
            [eVals,eSamps] = repair_broken_trigs(eSamps,eVals,settings);
        end

        windowTrigs.stim1 = eSamps(eVals == 65321); % STIM1 window trig (41)
        windowTrigs.dist = eSamps(eVals == 65322); % DIST window trig (42)
        windowTrigs.stim2 = eSamps(eVals == 65323); % STIM2 window trig (43)
        windowTrigs.stim2end = eSamps(eVals == 65303); % last STIM2 (23)

        % Define the trial structure -- no reconstruction necessary
        [trl,windowTrigs] = ft_trialfun_VAST_TFR(windowTrigs,settings, 0); 
    end

    windowTrigs.same = eSamps(eVals == 65341); % 1 keypress = "same"
    windowTrigs.diff = eSamps(eVals == 65342); % 2 keypress = "different"
    windowTrigs.sttrl = eSamps(eVals == 65380); % 100
    windowTrigs.endtrl = eSamps(eVals == 65381); % 101

    % More corrections for bad subjects
    if strcmp(settings.filenameNoExt, 'p6554_a')
        windowTrigs.sttrl = windowTrigs.sttrl(21:end);
        windowTrigs.endtrl(windowTrigs.endtrl < windowTrigs.sttrl(1)) = [];
        windowTrigs.same(windowTrigs.same < windowTrigs.sttrl(1)) = [];
        windowTrigs.diff(windowTrigs.diff < windowTrigs.sttrl(1)) = [];
    elseif strcmp(settings.filenameNoExt, 'p8057_a')
        offset500 = settings.ds2 * 0.5;
        stim2 = eSamps(eVals == 65300);
        stim2 = stim2(122) - offset500;
        windowTrigs.stim2 = [windowTrigs.stim2(1:121), stim2,...
            windowTrigs.stim2(122:end)];
    end

    % SAVE
    save([settings.filedirAnalyzed settings.filenameNoExt '_STEP4_vars'],...
        'eventTrigs', 'windowTrigs', 'trl', 'eSamps', 'eVals', '-v7.3');

end % Step 4


%% Step 5: Subtract condition-and-participant-specific ERPs

if STARTSTEP <= 5
    erp_file = strcat(erpdir, settings.filenameNoExt, '_avgERPs.mat');
    if exist(erp_file, 'file') == 2
        load(erp_file, 'avgERPs');
    else
        warning('It appears that you need to run VAST_preproc_ERP on this subject.')
        error('average ERPs have not yet been saved for this participant')
    end
    
    %%% --- Figure out to which trial each ERP event belongs
    eVals = bitand(eVals, 255);
    
    % Special corrections for one subject with weird triggers 
    if strcmp(settings.filenameNoExt, 'p6554_a')
        uniqTrigs = unique(eVals);
        uniqTrigs = uniqTrigs(~ismember(uniqTrigs, [0, 61, 62, 101, 255]));
        for tval = uniqTrigs
            findTrigs = find(eVals == tval);
            eVals(findTrigs(1:20)) = [];
            eSamps(findTrigs(1:20)) = [];
        end
        % do end trigs separately: trial 20 missing
        findTrigs = find(eVals == 101);
        eVals(findTrigs(1:19)) = [];
        eSamps(findTrigs(1:19)) = [];
        % get rid of any responses before the first start trial trig
        firstSt = find(eVals==100, 1);
        if firstSt > 1
            eSamps(1:firstSt - 1) = [];
            eVals(1:firstSt - 1) = [];
        end
    end

    trialAssign = assign_ERPs_to_trials(eSamps, eVals, settings.Ntrl);
    
    % Link conditions to trial numbers
    trlConds = get_trial_conditions(settings);
    
    keyTrigs = [{'enc_s1','enc_s2','enc_s3','enc_s4','int_s1','int_s2',...
        'int_s3','ret_s1','ret_s2','ret_s3','ret_s4'};{10, 11, 12, 13,...
        30, 31, 32, 20, 21, 22, 23}];
    load('dummy_indSs.mat', 'dummy_indSs')
    trl_t = dummy_indSs.time >= 0;
    len_t = sum(trl_t); % number of time points involved in each sub
    clear dummy_indSs
    nRawSamps = size(rawData.trial{1}, 2);
    badTrigTrials = zeros(1, settings.Ntrl); % triggers beyond raw data
    for t = 1:size(keyTrigs, 2)
        currType = keyTrigs{1,t};
        currTrigs = eSamps(eVals == keyTrigs{2,t});
        % to which trial number does each trigger belong?...
        trigTrials = trialAssign(eVals == keyTrigs{2,t});
        for i = 1:length(currTrigs)
            % Pull the average ERP corresponding to this trial and type
            trigCond = trlConds{trigTrials(i)};
            if contains(currType, 'int') && contains(trigCond, 'none')
                % skip subtraction -- no interfering ERPs
            else
                condFind = strcmp(avgERPs(1,:),trigCond);
                erp = avgERPs{2,condFind}{3,strcmp(avgERPs{2,condFind}(1,:),...
                    currType)};
                erp = erp(:,trl_t); % cut off pre-event baseline
                % Perform the subtraction
                if currTrigs(i) + len_t - 1 > nRawSamps
                    badTrigTrials(trigTrials(i)) = 1; % trig exceeds data
                else
                    rawChunk = rawData.trial{1}(:,...
                        currTrigs(i):currTrigs(i) + len_t - 1);
                    rawData.trial{1}(:,currTrigs(i):currTrigs(i) +...
                        len_t - 1) = rawChunk - erp;
                end
            end
        end % individual trig loop
    end % trig type loop
    
    badTrigTrials = logical(badTrigTrials);
    
    disp(['Saving: ',settings.filedirAnalyzed, settings.filenameNoExt, '_5_ds_filt_ica_ERPsub'])
    save([settings.filedirAnalyzed settings.filenameNoExt '_STEP5_vars'],...
        'badTrigTrials', 'trlConds', '-v7.3');
    save([settings.filedirAnalyzed settings.filenameNoExt '_5_ds_filt_ica_ERPsub'],'rawData','-v7.3');
    
end % Step 5


%% Step 6: Epoch into FULL TRIALS (with padding)

if STARTSTEP <= 6
    
    % Remove trials that cut out any of the triggers
    trl_cln1 = trl;
    trl_cln1(badTrigTrials, :) = [];
    
    % Try to solve other padding issues
    Nrawsamps = size(rawData.trial{1}, 2);
    lastTrlPadProblem = false;
    % start trial triggers being beyond data is highly unusual
    if any(trl_cln1(:,1) < 0) || any(trl_cln1(:,1) > Nrawsamps)
        error("Investigate issues with trial start definitions.")
    % so is having more than one trial with messed up end trial trigs
    elseif sum(trl_cln1(:,2) > Nrawsamps) > 1
        error("Investigate issues with trial end definitions.")
    % can't fix end trial trig issue other than the last trial 
    elseif sum(trl_cln1(:,2) > Nrawsamps) == 1 &&...
            find(trl_cln1(:,2) > Nrawsamps) ~= size(trl, 1)
        error("Investigate issues with trial end definitions.")
    % otherwise, a last trial padding problem can be fixed
    elseif sum(trl_cln1(:,2) > Nrawsamps) == 1
        Nsamps2makeup = trl_cln1(end,2) - Nrawsamps;
        if Nsamps2makeup <= 2 * settings.ds2 % just padding, can correct 
            warning("Read padding will be augmented with manual mirroring")
            % for now, just mark it and allow the data to be epoched
            lastTrlPadProblem = true;
            trl_cln_adj = trl_cln1;
            trl_cln_adj(end,2) = Nrawsamps;
        else % too many missing samples
            error("Investigate issues with trial end definitions.")
        end
    end

    % Note: if issues arise that require manual padding (i.e. by time
    % reversal), see checkout this function from commit #83e077a

    % Epoch based on trl
    cfg = [];
    if lastTrlPadProblem
        cfg.trl = trl_cln_adj;
    else
        cfg.trl = trl_cln1;
    end
    
    rawData = ft_redefinetrial(cfg, rawData);
    
    % If there was a padding problem, fix it manually
    if lastTrlPadProblem
        lastdata = rawData.trial{end};
        manpad = fliplr(lastdata(:, end - Nsamps2makeup + 1:end));
        rawData.trial{end} = [lastdata, manpad];
        % repair the timebase as well
        lasttime = rawData.time{end};
        timestep = lasttime(2) - lasttime(1);
        newtimes = lasttime(end) + timestep : timestep : ...
                   lasttime(end) + (timestep*Nsamps2makeup);
        rawData.time{end} = [rawData.time{end}, newtimes];
    end
    
    % Remove the missing trials from trigger structs
    windowTrigs_cln.baseline_st = windowTrigs.baseline_st(badTrigTrials==0);
    windowTrigs_cln.stim1 = windowTrigs.stim1(badTrigTrials==0);
    windowTrigs_cln.dist = windowTrigs.dist(badTrigTrials==0);
    windowTrigs_cln.stim2 = windowTrigs.stim2(badTrigTrials==0);
    windowTrigs_cln.stim2end = windowTrigs.stim2end(badTrigTrials==0);

    disp(['Saving: ',settings.filedirAnalyzed, settings.filenameNoExt, '_6_ds_filt_ica__ERPsub_ep'])
    save([settings.filedirAnalyzed settings.filenameNoExt '_STEP6_vars'],...
        'trl_cln1', 'windowTrigs_cln', '-v7.3');
    save([settings.filedirAnalyzed settings.filenameNoExt '_6_ds_filt_ica_ERPsub_ep'],'rawData','-v7.3');

end % Step 6


%% Step 7: Peak-to-peak artifact detection in individual stim windows

if STARTSTEP <= 7
    Nepoch = length(rawData.trial);
    [badEpochs.bl, badEpochs.enc, badEpochs.int, badEpochs.ret] =...
        deal(false(1, Nepoch));
    
    [rel.enc_st, rel.int_st, rel.ret_st, rel.ret_end] =...
        deal(zeros(1, Nepoch));
    rel.baseline = 1;
    
    for rTr = 1:Nepoch
        currdata = rawData.trial{rTr};
        % ignore the first and last 2 sec -- just padding
        samps2sec = settings.ds2 * 2;
        currdata(:,1:samps2sec) = []; % start at windowTrigs.baseline_st
        currdata(:,end-samps2sec+1:end) = [];

        % Find regions that exceed the rejection threshold
        currdata = (currdata > settings.p2pThresh_max) |...
                   (currdata < settings.p2pThresh_min);
        % Use trig vals to get the relative positions of each window
        abs_baseline = windowTrigs_cln.baseline_st(rTr);
        rel.enc_st(rTr) = windowTrigs_cln.stim1(rTr) - abs_baseline + 1;
        rel.int_st(rTr) = windowTrigs_cln.dist(rTr) - abs_baseline + 1;
        rel.ret_st(rTr) = windowTrigs_cln.stim2(rTr) - abs_baseline + 1;
        rel.ret_end(rTr) = windowTrigs_cln.stim2end(rTr) - abs_baseline + 1;

        % Check for bad epochs...
        if sum(any(currdata(:,rel.baseline : rel.enc_st(rTr) - 1)))
            badEpochs.bl(rTr) = true; % baseline window
        end
        % if baseline is bad, can't work with any of the other windows
        if badEpochs.bl(rTr) == true
            [badEpochs.enc(rTr), badEpochs.int(rTr), badEpochs.ret(rTr)] =...
                deal(true);
        else % baseline ok; check the other windows individually 
            if sum(any(currdata(:, rel.enc_st(rTr) : rel.int_st(rTr) - 1)))
                badEpochs.enc(rTr) = true; % encoding window
            end
            if sum(any(currdata(:, rel.int_st(rTr) : rel.ret_st(rTr) - 1)))
                badEpochs.int(rTr) = true; % interfering/hold window
            end
            if sum(any(currdata(:, rel.ret_st(rTr) : rel.ret_end(rTr))))
                badEpochs.ret(rTr) = true; % retrieval window
            end
        end
    end

    fprintf('Automatic Rejection Rate, ENC: %0.02f \n',...
        sum(badEpochs.enc)/length(badEpochs.enc))
    fprintf('Automatic Rejection Rate, INT: %0.02f \n',...
        sum(badEpochs.int)/length(badEpochs.int))
    fprintf('Automatic Rejection Rate, RET: %0.02f \n',...
        sum(badEpochs.ret)/length(badEpochs.ret))
    
    save([settings.filedirAnalyzed settings.filenameNoExt '_STEP7_vars'],...
        'badEpochs', 'rel', '-v7.3');

end % Step 7


%% Step 8: Short-time Fourier Transform

if STARTSTEP <= 8

    % Get the ACTUAL trial lengths -- baseline to stim2end -- from
    % windowTrigs
    trlLengths = zeros(1, length(rawData.trial));
    for rTr = 1:length(rawData.trial)
        currBL_st = windowTrigs_cln.baseline_st(rTr);
        currRET_end = windowTrigs_cln.stim2end(rTr);
        trlLengths(rTr) = length(currBL_st:currRET_end) / settings.ds2;
    end

    maxTrl = max(trlLengths); % sec; longest trial
    tSpacing = 0.1; % sec, time spacing of resulting TFRs (100 ms)
    numTpts = floor(maxTrl/tSpacing); 
    
    % Need even more padding -- try copying, mirroring, and appending the
    % first and last 3 sec of each trial
    
    lengthExtraPad = 3; % sec
    for t = 1:length(rawData.trial)
        sampsRep_beg = fliplr(rawData.trial{t}(:, rawData.time{t} <...
            lengthExtraPad));
        sampsRep_end = fliplr(rawData.trial{t}(:,...
            rawData.time{t} > rawData.time{t}(end) - lengthExtraPad));
        rawData.trial{t} = [sampsRep_beg, rawData.trial{t}, sampsRep_end];
        % Reconstruct a timebase
        Nsamps = length(sampsRep_beg);
        t_start = rawData.time{t}(1:Nsamps);
        t_adj = rawData.time{t} + lengthExtraPad;
        t_adj_end = t_adj(end-Nsamps+1:end) + lengthExtraPad;
        rawData.time{t} = [t_start, t_adj, t_adj_end];
    end
    
    % Compute TFR on each complete trial -- WAVELET VERSION
    ftcfg = [];
    ftcfg.channel = 'all';
    ftcfg.method = 'wavelet';
    ftcfg.output = 'pow';
    ftcfg.width = 5; % higher -> freq res inc, temporal res dec, & vice versa
%     ftcfg.foi = 1:1:40;
    ftcfg.foi = 1:1:80;
    % 2 sec padding is read in. [lengthExtraPad] sec padding added above.
    ftcfg.toi = linspace(2 + lengthExtraPad,...
        2 + lengthExtraPad + maxTrl, numTpts); % times of interest
    ftcfg.padtype = 'mirror';
    ftcfg.pad = 'nextpow2';
    ftcfg.keeptrials = 'yes'; % don't average at this step
    
    TFR = ft_freqanalysis(ftcfg, rawData);
    TFR.time = TFR.time - (2 + lengthExtraPad); % account for padding

    disp(['Saving: ',settings.filedirAnalyzed, settings.filenameNoExt, '_TFR'])
    save([settings.filedirAnalyzed settings.filenameNoExt '_TFR.mat'],'TFR','-v7.3');
    save([settings.filedirAnalyzed settings.filenameNoExt '_STEP8_vars'],...
        'trlLengths','-v7.3');

end % Step 8


%% Step 9: Split TFR into responses in each stim window (ENC, INT, RET)
% NOTE: Pre-trial baseline attached to each window

if STARTSTEP <= 9
    
    % Calculate the grand average baseline
    blInds = TFR.time <= 1; % after 1 sec baseline
    BL = TFR.powspctrm(~badEpochs.bl,:,:,blInds); % don't avg bad epochs
    BL = squeeze(mean(BL, 1));
    
    % Calculate condition-specific baselines
    trlConds_noTrigErr = trlConds(~badTrigTrials);
    cNames = unique(trlConds_noTrigErr);
    cNames = cNames([5,4,6,2,1,3]); % put names in preferred order
    BL_cSpec = [cNames; cell(1,length(cNames))];
    for c = 1:size(BL_cSpec, 2)
        currC = contains(trlConds_noTrigErr, BL_cSpec{1,c});
        currC(badEpochs.bl) = false;
        temp = TFR.powspctrm(currC,:,:,blInds);
        BL_cSpec{2,c} = squeeze(mean(temp, 1));
    end
    clear temp
    
    Ntrl = size(trl_cln1, 1);
    
    % Get trial-relative triggers in units of time instead of samples
    rel_t = rel;
    clear rel
    rel_t.enc_st = rel_t.enc_st ./ settings.ds2;
    rel_t.int_st = rel_t.int_st ./ settings.ds2;
    rel_t.ret_st = rel_t.ret_st ./ settings.ds2;
    rel_t.ret_end = rel_t.ret_end ./ settings.ds2;

    % Get the samples in TFR time for each window...
    maxINT = 4.5; % sec, max length of INT window
    % NOTE re. maxINT: In the real experiment, the actual time between the
    % end of STIM1 and the start of STIM2 (full INT window) is 5.5 sec.
    % However, the trialfun defines the start of INT as 500ms after the
    % last STIM1 stimulus and the end of INT as 500ms before the first
    % STIM2 stimulus. Thus, the maximum achievable INT window length (as 
    % defined here) is 4.5 sec.
    maxINT_samps = find(TFR.time <= maxINT, 1, 'last');
    % relative baseline times don't change...
    tfrsamps.bl = find(TFR.time < rel_t.enc_st(1), 1, 'last');
    % Loop over other window markers
    [tfrsamps.enc, tfrsamps.int, tfrsamps.ret] = deal(zeros(Ntrl, 2));
    for t = 1:Ntrl
        tfrsamps.enc(t,1) = find(TFR.time >= rel_t.enc_st(t), 1);
        tfrsamps.enc(t,2) = find(TFR.time < rel_t.int_st(t), 1, 'last');
        tfrsamps.int(t,1) = find(TFR.time >= rel_t.int_st(t), 1);
        tfrsamps.int(t,2) = find(TFR.time < rel_t.ret_st(t), 1, 'last');
        tfrsamps.ret(t,1) = find(TFR.time >= rel_t.ret_st(t), 1);
        tfrsamps.ret(t,2) = find(TFR.time <= rel_t.ret_end(t), 1, 'last');
        
        % Truncate the INT window if it ran too long
        if length(tfrsamps.int(t,1):tfrsamps.int(t,2)) > maxINT_samps
            tfrsamps.int(t,2) = tfrsamps.int(t,1) + (maxINT_samps - 1);
        end
    end
    
    % Preallocate the window TFRs
    TFRcopy = struct();
    TFRcopy.label = TFR.label;
    TFRcopy.dimord = TFR.dimord;
    TFRcopy.freq = TFR.freq;
    TFRcopy.cumtapcnt = TFR.cumtapcnt;
    TFRcopy.elec = TFR.elec;
    TFRcopy.cfg = TFR.cfg;
    [tfr_enc, tfr_int, tfr_ret] = deal(TFRcopy);
    
    sizeTFR = size(TFR.powspctrm);
    tfr_enc.powspctrm = nan([sizeTFR(1:3), (tfrsamps.bl +...
        max(tfrsamps.enc(:,2) - tfrsamps.enc(:,1)) + 1)]);
    tfr_enc.time = TFR.time(1:size(tfr_enc.powspctrm, 4));
    tfr_int.powspctrm = nan([sizeTFR(1:3), (tfrsamps.bl +...
        max(tfrsamps.int(:,2) - tfrsamps.int(:,1)) + 1)]);
    tfr_int.time = TFR.time(1:size(tfr_int.powspctrm, 4));
    tfr_ret.powspctrm = nan([sizeTFR(1:3), (tfrsamps.bl +...
    	max(tfrsamps.ret(:,2) - tfrsamps.ret(:,1)) + 1)]);
    tfr_ret.time = TFR.time(1:size(tfr_ret.powspctrm, 4));
                        
    % Cut up the TFR into windows
    for t = 1:Ntrl
        currdata = TFR.powspctrm(t,:,:,:);
        % Plug in the pre-trial baseline to each window
        currdata_bl = currdata(1,:,:,1:tfrsamps.bl);
        [tfr_enc.powspctrm(t,:,:,1:tfrsamps.bl),...
            tfr_int.powspctrm(t,:,:,1:tfrsamps.bl),...
            tfr_ret.powspctrm(t,:,:,1:tfrsamps.bl)] = deal(currdata_bl);
        % ENC window
        tfr_enc.powspctrm(t,:,:,tfrsamps.enc(t,1):tfrsamps.enc(t,2)) =...
            currdata(1,:,:,tfrsamps.enc(t,1):tfrsamps.enc(t,2));
        % ...for other two windows, need to adjust tfrsamps to occur just
        %    after the window baseline...
        % INT window
        adjint = tfrsamps.int(t,:) - (tfrsamps.int(t,1) - tfrsamps.bl - 1);
        tfr_int.powspctrm(t,:,:,adjint(1):adjint(2)) =...
            currdata(1,:,:,tfrsamps.int(t,1):tfrsamps.int(t,2));
        % RET window
        adjret = tfrsamps.ret(t,:) - (tfrsamps.ret(t,1) - tfrsamps.bl - 1);
        tfr_ret.powspctrm(t,:,:,adjret(1):adjret(2)) =...
            currdata(1,:,:,tfrsamps.ret(t,1):tfrsamps.ret(t,2));
    end
    
    % Get rid of previously marked bad epochs
    tfr_enc.powspctrm(badEpochs.enc,:,:,:) = [];
    tfr_enc.cumtapcnt(badEpochs.enc,:) = [];
    tfr_int.powspctrm(badEpochs.int,:,:,:) = [];
    tfr_int.cumtapcnt(badEpochs.int,:) = [];
    tfr_ret.powspctrm(badEpochs.ret,:,:,:) = [];
    tfr_ret.cumtapcnt(badEpochs.ret,:) = [];
    
    % Add badTrigTrials to badEpochs for ease of downstream analysis
    if any(badTrigTrials)
        whereBadTrigs = find(badTrigTrials);
        for tloc = whereBadTrigs
            % Note: on last iteration, tloc may exceed length of badEpochs.
            % This is fine, last item will just be 1x0 empty vector.
            badEpochs.bl = [badEpochs.bl(1:tloc-1), true,...
                badEpochs.bl(tloc:end)];
            badEpochs.enc = [badEpochs.enc(1:tloc-1), true,...
                badEpochs.enc(tloc:end)];
            badEpochs.int = [badEpochs.int(1:tloc-1), true,...
                badEpochs.int(tloc:end)];
            badEpochs.ret = [badEpochs.ret(1:tloc-1), true,...
                badEpochs.ret(tloc:end)];
        end
    end
    
    disp(['Saving: ',settings.filedirAnalyzed, settings.filenameNoExt, '_TFR_bywin'])
    save([settings.filedirAnalyzed settings.filenameNoExt '_TFR_bywin'],...
        'tfr_enc', 'tfr_int', 'tfr_ret', '-v7.3');
    save([settings.filedirAnalyzed settings.filenameNoExt '_avg_TFR_BL'],...
        'BL', '-v7.3');
    save([settings.filedirAnalyzed settings.filenameNoExt '_cSpec_TFR_BL'],...
        'BL_cSpec', '-v7.3');
    save([settings.filedirAnalyzed settings.filenameNoExt '_STEP9_vars'],...
        'rel_t', 'tfrsamps', 'badEpochs', '-v7.3');

end % Step 9
    


%%%________________________________________________________________________
%% HELPER FUNCTIONS
%%%________________________________________________________________________


%% Assign a trial number to each trigger

function trialAssign = assign_ERPs_to_trials(eSamps, eVals, Ntrl)

if sum(eVals == 100) ~= Ntrl
    error("Missing at least one Start Trial trigger. Cannot perform condition assignment.") 
end
if sum(eVals == 101) ~= Ntrl
    warning('Missing at least one End Trial trigger.')
    warning('Attempting trial assign with only Start Trial triggers')
    
    trialAssign = zeros(size(eSamps));
    startSamps = eSamps(eVals == 100);
    for t = 1:Ntrl
        if t == Ntrl
            trialAssign(eSamps >= startSamps(t)) = t;
        else
            trialAssign(eSamps >= startSamps(t) &...
                eSamps < startSamps(t+1)) = t;
        end
        
    end
else    
    trialAssign = zeros(size(eSamps));
    startSamps = eSamps(eVals == 100);
    lastSamps = eSamps(eVals == 101);
    for t = 1:Ntrl
         trialAssign(eSamps >= startSamps(t) & eSamps <= lastSamps(t)) = t;
    end
end



%% Get trial conditions from experiment cfg

function trlConds = get_trial_conditions(settings)

% Load the corresponding experiment cfg
allfiles = dir(settings.filedirRaw);
allfiles = extractfield(allfiles, 'name');
allfiles = allfiles(3:end);
isSubj = contains(allfiles, settings.filenameNoExt);
allfiles = allfiles(isSubj);
isMat = contains(allfiles, ".mat");
if sum(isMat) > 1
    error("more than one experiment .mat file -- not sure which to load")
end
behavfile = allfiles{isMat};
expcfg = load(strcat(settings.filedirRaw, behavfile), "cfg");
expcfg = expcfg.cfg;

% figure out condition assignment 
trlConds = cell(1, expcfg.nBlocks * expcfg.nTrials);
for b = 1:expcfg.nBlocks
    m = expcfg.m1{b};
    d = expcfg.d{b};
    int = expcfg.dist{b};
    currcond = strcat(m, d, '_', int);
    inds = ((b-1)*expcfg.nTrials) + 1 : b * expcfg.nTrials;
    trlConds(inds) = {currcond};
end


